package com.ywu.standard.gateway.handler.header;

import com.ywu.boot.context.YwuContext;
import com.ywu.boot.context.YwuContextHolder;
import com.ywu.boot.context.YwuContextRequest;
import com.ywu.standard.gateway.handler.header.impl.HeaderHandlerInterface;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.annotation.Order;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;

import java.util.Objects;
import java.util.UUID;

/**
 * @ClassName HeaderTraceHandler
 * @Description TODO
 * @Author GroundDemo
 * @Date 2024/5/26 11:04
 * @Version 1.0
 **/
@Component
@Slf4j
@Order(1)
public class HeaderTraceHandler implements HeaderHandlerInterface {
    private static final String traceIdHeaderKey = "X-Pix-Ywu-Trace-Id";

    private static final String spanIdHeaderKey = "X-Pix-Ywu-Span-Id";

    /**
    * @Author GroundDemo
    * @Description 请求头trace信息添加
    * @Date 11:11 2024/5/26
    * @Param [exchange]
    * @return org.springframework.web.server.ServerWebExchange
    **/
    @Override
    public Object handle(ServerWebExchange exchange) {
        log.debug("HeaderTraceHandler begin");
        String id = UUID.randomUUID().toString().replace("-", "");
        ServerHttpRequest mutatedTraceRequest = exchange.getRequest().mutate().header(traceIdHeaderKey, id).build();
        ServerWebExchange traceExchange = exchange.mutate().request(mutatedTraceRequest).build();
        ServerHttpRequest mutatedSpanRequest = traceExchange.getRequest().mutate().header(spanIdHeaderKey, id).build();
        YwuContext currentContext = YwuContextHolder.getCurrentContext();
        if (Objects.isNull(currentContext)) {
            currentContext = new YwuContext();
        }
        currentContext.setYwuContextRequest(new YwuContextRequest().setTraceId(id).setSpanId(id));
        YwuContextHolder.setCurrentContext(currentContext);
        return exchange.mutate().request(mutatedSpanRequest).build();
    }
}
